const readonlyEmptyArray = Object.freeze([])

const ArrayPrototype: any = Array.prototype
const SetPrototype: any = Set.prototype

ArrayPrototype.delete = function (item: any) {
  const idx = this.indexOf(item)
  if (idx >= 0) {
    this.splice(idx, 1)
    return true
  }
  return false
}

ArrayPrototype.clear = function () {
  this.splice(0)
}

SetPrototype.push = function (...items: any[]) {
  for (const value of items) {
    this.add(value)
  }
}

SetPrototype.includes = SetPrototype.has

Object.defineProperty(SetPrototype, 'length', {
  get(this: any) {
    return this.size
  },
})

const _Map = Map as new (...args: any[]) => any
const _WeakMap = WeakMap as new (...args: any[]) => any

export class MultiMap<K, V> extends _Map {
  private readonly __: boolean

  constructor(entries?: Iterable<(readonly [K, V])> | null, unique = false) {
    super()
    this.__ = unique

    if (entries) {
      for (const [key, value] of entries) {
        this.set(key, value)
      }
    }
  }

  get(key: K): Iterable<V> {
    return super.get(key) ?? readonlyEmptyArray
  }

  set(key: K, ...values: V[]) {
    let storage_values = super.get(key)
    if(!storage_values) {
      storage_values = this.__ ? new Set() : []
      super.set(key, storage_values)
    }
    storage_values.push(...values)
  }

  has(key: K, value?: V): boolean {
    if (value !== undefined) {
      return !!super.get(key)?.includes(value)
    } else {
      return (super.get(key)?.length ?? 0) > 0
    }
  }

  delete(key: K, value?: V): boolean {
    const storage_values = super.get(key)
    if (!storage_values) return false

    let result: boolean

    if (value !== undefined) {
      result = storage_values.delete(value)
      if (storage_values.length === 0) {
        super.delete(key)
      }
    } else {
      result = storage_values.length > 0
      super.delete(key)
    }

    return result
  }

  clear(): void {
    super.clear()
  }

  keys(): IterableIterator<K>  {
    return super.keys()
  }

  *values(): IterableIterator<V> {
    for (const values of super.values()) {
      for (const value of values) {
        yield value
      }
    }
  }

  *entries(): IterableIterator<[K, V]> {
    for(const [key, values] of super.entries()) {
      for (const value of values) {
        yield [key, value]
      }
    }
  }

  *[Symbol.iterator](): IterableIterator<[K, V]> {
    for(const [key, values] of super[Symbol.iterator]()) {
      for (const value of values) {
        yield [key, value]
      }
    }
  }

  get size(): number {
    return super.size
  }

  get count(): number {
    let i = 0
    for (const value of this.values()) {
      ++i
    }
    return i
  }

  get [Symbol.toStringTag]() {
    return 'MultiMap'
  }
}

export class MultiWeakMap<K extends object, V> extends _WeakMap {
  private readonly __: boolean

  constructor(entries?: Iterable<(readonly [K, V])> | null, unique = false) {
    super()
    this.__ = unique

    if (entries) {
      for (const [key, value] of entries) {
        this.set(key, value)
      }
    }
  }

  get(key: K): Iterable<V> {
    return super.get(key) ?? readonlyEmptyArray
  }

  set(key: K, ...values: V[]) {
    let storage_values = super.get(key)
    if(!storage_values) {
      storage_values = this.__ ? new Set() : []
      super.set(key, storage_values)
    }
    storage_values.push(...values)
  }

  has(key: K, value?: V): boolean {
    if (value !== undefined) {
      return !!super.get(key)?.includes(value)
    } else {
      return (super.get(key)?.length ?? 0) > 0
    }
  }

  delete(key: K, value?: V): boolean {
    const storage_values = super.get(key)
    if (!storage_values) return false

    let result: boolean

    if (value !== undefined) {
      result = storage_values.delete(value)
      if (storage_values.length === 0) {
        super.delete(key)
      }
    } else {
      result = storage_values.length > 0
      super.delete(key)
    }
    
    return result
  }

  get [Symbol.toStringTag]() {
    return 'MultiWeakMap'
  }
}

declare global {

  interface Array<T> {
    delete(item: T): boolean
    clear(): void
  }

  interface Set<T> {
    readonly length: number
    push(...items: T[]): void
    includes(item: T): boolean
  }
}
